package main

import (
	"MASSTEST/poc"
	"MASSTEST/poc/pocutil"
	"bytes"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	//"net"
	"os"
	"sort"
	//massdb_v1_test "MASSTEST/poc/engine/massdb/massdb.v1"
	"MASSTEST/poc/engine"
	massdb_v1 "MASSTEST/poc/engine/massdb/massdb.v1"
	"MASSTEST/pocec"
	"encoding/gob"

	"github.com/urfave/cli/v2"
)

//Proof type for encoding
type Proof struct {
	X         []byte
	Xp        []byte
	Success   bool
	Challenge pocutil.Hash
}

type SigProof struct {
	ProofSer []byte
	SigSer   []byte
	Bl       int
}

//// 仅持有1/2大小hashmap有99.8%概率无法通过
//var (
//	challengeNumber  = 300
//	minSuccessNumber = 70 // false negative rate = 0.0001
//)
//

// 上面参数会导致命令行参数在Windows下超过限制长度，这组参数用于测试
var (
	challengeNumber  = 30
	minSuccessNumber = 5
)

func main() {
	app := &cli.App{
		Commands: []*cli.Command{
			{
				Name: "gen",
				Subcommands: []*cli.Command{
					{
						Name:   "key",
						Usage:  "generate a pair of pk&sk",
						Action: genKey,
					},
					{
						Name:   "challenge",
						Usage:  "generate n chanllenges to check the client",
						Action: genChallenge,
						Flags: []cli.Flag{
							&cli.Int64Flag{
								Name:  "n",
								Value: int64(challengeNumber),
								Usage: "number of chanllenges",
							},
							&cli.StringFlag{
								Name:    "seed",
								Aliases: []string{"s"},
								Value:   "MASS",
								Usage:   "seed to generate challenge",
							},
						},
					},
					{
						Name:  "proof",
						Usage: "generate proof of chanllenges",
						Flags: []cli.Flag{
							&cli.IntFlag{
								Name:    "bl",
								Aliases: []string{"b"},
								Value:   24,
								Usage:   "block size",
							},
							&cli.StringFlag{
								Name:    "challenge",
								Aliases: []string{"c"},
								Usage:   "chanllenge to proof",
							},
							&cli.StringFlag{
								Name:    "sk",
								Aliases: []string{"k"},
								Usage:   "private key of hashmap",
							},
							&cli.StringFlag{
								Name:    "rootpath",
								Aliases: []string{"p"},
								Value:   "testdata",
								Usage:   "path of hashmap data",
							},
						},
						Action: genProof,
					},
				},
			},
			{
				Name: "verify",
				Subcommands: []*cli.Command{
					{
						Name:  "proof",
						Usage: "verify proof",
						Flags: []cli.Flag{
							&cli.StringFlag{
								Name:    "proof",
								Aliases: []string{"p"},
								Usage:   "output of genproof",
							},
							&cli.StringFlag{
								Name:    "pk",
								Aliases: []string{"k"},
								Usage:   "public key of hashmap",
							},
							&cli.StringFlag{
								Name:    "seed",
								Aliases: []string{"s"},
								Value:   "MASS",
								Usage:   "seed of challenge",
							},
						},
						Action: verifyProof,
					},
					{
						Name:  "info",
						Usage: "verify info file",
						Flags: []cli.Flag{
							&cli.StringFlag{
								Name:    "md5",
								Aliases: []string{"m"},
								Usage:   "md5 of hashmap",
							},
							&cli.StringFlag{
								Name:    "sha1",
								Aliases: []string{"s"},
								Usage:   "sha1 of hashmap",
							},
							&cli.StringFlag{
								Name:    "url",
								Aliases: []string{"u"},
								Value:   "",
								Usage:   "url of callback api",
							},
							&cli.StringFlag{
								Name:    "pk",
								Aliases: []string{"k"},
								Usage:   "public key of hashmap",
							},
							&cli.StringFlag{
								Name:    "sig",
								Aliases: []string{"x"},
								Value:   "",
								Usage:   "signature of info",
							},
						},
						Action: verifyInfo,
					},
				},
			},
			{
				Name:  "init",
				Usage: "Init hashmap",
				Flags: []cli.Flag{
					&cli.IntFlag{
						Name:    "bl",
						Aliases: []string{"b"},
						Value:   24,
						Usage:   "block size",
					},
					&cli.StringFlag{
						Name:    "rootpath",
						Aliases: []string{"p"},
						Value:   "testdata",
						Usage:   "specify the path of hashmap data",
					},
					&cli.StringFlag{
						Name:    "url",
						Aliases: []string{"u"},
						Value:   "",
						Usage:   "url of callback api",
					},
					&cli.StringFlag{
						Name:    "sk",
						Aliases: []string{"k"},
						Usage:   "private key of hashmap",
					},
				},
				Action: plot,
			},
		},
	}

	sort.Sort(cli.FlagsByName(app.Flags))
	sort.Sort(cli.CommandsByName(app.Commands))

	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func genKey(_ *cli.Context) error {
	sk, err := pocec.NewPrivateKey(pocec.S256())
	if err != nil {
		fmt.Println(err)
		return err
	}
	pk := sk.PubKey()

	fmt.Printf("publicKey:\n%s\n", hex.EncodeToString(pk.SerializeCompressed()))
	fmt.Printf("privateKey:\n%s\n", hex.EncodeToString(sk.Serialize()))
	return nil
}

func genChallenge(c *cli.Context) error {
	var n = c.Int("n")
	var seed = c.String("seed")
	var chanllengeCollection []pocutil.Hash
	var challengeBuffer bytes.Buffer

	challengeEncoder := gob.NewEncoder(&challengeBuffer)

	challenge := pocutil.SHA256([]byte(seed))
	for i := 0; i < n; i++ {
		challenge = pocutil.SHA256(challenge[:])
		chanllengeCollection = append(chanllengeCollection, challenge)
	}
	err := challengeEncoder.Encode(&chanllengeCollection)
	if err != nil {
		return err
	}
	fmt.Printf("Challenge:\n%s\n", hex.EncodeToString(challengeBuffer.Bytes()))
	return nil
}

func plot(c *cli.Context) error {
	type config struct {
		Md5  string `json:"md5"`
		Sha1 string `json:"sha1"`
		Url  string `json:"callback"`
		Pk   string `json:"publicKey"`
		Sig  string `json:"signature"`
	}
	var url = c.String("url")
	var bl = c.Int("bl")
	var rootPath = c.String("rootpath")
	var sSk = c.String("sk")
	var configBuffer bytes.Buffer

	sk, pk, err := privateKeyFromString(sSk)
	if err != nil {
		return err
	}

	//t1 := time.Now()

	mdb, err := massdb_v1.NewMassDBV1(rootPath, engine.UnknownOrdinal, pk, bl)
	if err != nil {
		return err
	}
	result := mdb.Plot()
	if err = <-result; err != nil {
		return err
	}

	// 取得路径
	path := getPath(rootPath, engine.UnknownOrdinal, pk, bl)
	pathB := path + ".massdb"
	pathC := path + ".info"
	pHashMapFile, err := os.Open(pathB)
	defer pHashMapFile.Close()
	if err != nil {
		return err
	}

	// 计算hash
	md5Hash := md5.New()
	io.Copy(md5Hash, pHashMapFile)
	bMd5 := md5Hash.Sum(nil)
	sha1Hash := sha1.New()
	io.Copy(sha1Hash, pHashMapFile)
	bSha1 := sha1.Sum(nil)

	// 签名
	configBuffer.Write(bMd5[:])
	configBuffer.Write(bSha1[:])
	configBuffer.Write([]byte(url))
	configHash := sha256.Sum256(configBuffer.Bytes())
	s, err := sk.Sign(configHash[:])
	if err != nil {
		return err
	}

	// 生成json
	configJSON := &config{
		Url:  url,
		Md5:  hex.EncodeToString(bMd5[:]),
		Sha1: hex.EncodeToString(bSha1[:]),
		Sig:  hex.EncodeToString(s.Serialize()),
		Pk:   hex.EncodeToString(pk.SerializeCompressed()),
	}
	bConfigJSON, err := json.Marshal(configJSON)
	if err != nil {
		return err
	}
	pConfigFile, err := os.Create(pathC)
	defer pConfigFile.Close()
	if err != nil {
		return err
	}
	pConfigFile.Write(bConfigJSON)
	fmt.Printf("Path of info file:\n%s\n", pathC)
	//t2 := time.Now()
	//fmt.Println(t2.Sub(t1))
	return nil
}

func verifyProof(c *cli.Context) error {
	var sSigProof = c.String("proof")
	var targetSeed = c.String("seed")
	var sPk = c.String("pk")
	var sigProof = SigProof{}
	var bSigProof []byte
	var proofCollection []Proof

	bSigProof, err := hex.DecodeString(sSigProof)
	if err != nil {
		err = fmt.Errorf("Failed to decode the proof\n")
		return err
	}
	sigProofDecoder := gob.NewDecoder(bytes.NewReader(bSigProof))
	if err := sigProofDecoder.Decode(&sigProof); err != nil {
		return err
	}

	// 验证签名
	bPreProof := sigProof.ProofSer
	pk, err := pubKeyFromString(sPk)
	if err != nil {
		return err
	}
	proofHash := sha256.Sum256(bPreProof)
	sig, err := pocec.ParseSignature(sigProof.SigSer, pocec.S256())
	if err != nil {
		return err
	}
	if !sig.Verify(proofHash[:], pk) {
		err = fmt.Errorf("Failed to verify the signature\n")
		return err
	}

	proofDecoder := gob.NewDecoder(bytes.NewReader(bPreProof))
	if err := proofDecoder.Decode(&proofCollection); err != nil {
		return err
	}
	pkHash := pocutil.PubKeyHash(pk)

	// check每个proof
	targetHash := pocutil.SHA256([]byte(targetSeed))
	successCount := 0
	for _, proof := range proofCollection {
		// 与seed生成的期望challenge比较 防止伪造seed
		targetHash = pocutil.SHA256(targetHash[:])
		challengeHash := pocutil.Hash{}
		copy(challengeHash[:], proof.Challenge[:])
		if challengeHash != targetHash {
			err = fmt.Errorf("Wrong Seed\n")
			return err
		}

		if proof.Success {
			massProof := &poc.Proof{X: proof.X, XPrime: proof.Xp, BitLength: sigProof.Bl}
			if err := poc.VerifyProof(massProof, pkHash, proof.Challenge); err == nil {
				successCount += 1
			}
		}
	}

	if successCount >= minSuccessNumber {
		fmt.Println(true)
	} else {
		fmt.Println(false)
	}
	return nil
}

func verifyInfo(c *cli.Context) error {
	var sMd5 = c.String("md5")
	var sSha1 = c.String("sha1")
	var sPk = c.String("pk")
	var sSig = c.String("sig")
	var url = c.String("url")
	var configBuffer bytes.Buffer

	pk, err := pubKeyFromString(sPk)
	if err != nil {
		return err
	}
	sig, err := signatureFromString(sSig)
	if err != nil {
		return err
	}

	bMd5, err := hex.DecodeString(sMd5)
	if err != nil {
		return err
	}
	bSha1, err := hex.DecodeString(sSha1)
	if err != nil {
		return err
	}

	configBuffer.Write(bMd5)
	configBuffer.Write(bSha1)
	configBuffer.Write([]byte(url))
	configHash := sha256.Sum256(configBuffer.Bytes())
	if sig.Verify(configHash[:], pk) {
		fmt.Println(true)
	} else {
		fmt.Println(false)
	}

	return nil
}

func genProof(c *cli.Context) error {
	var bl = c.Int("bl")
	var sSk = c.String("sk")
	var rootPath = c.String("rootpath")
	var sChallenge = c.String("challenge")
	var proofBuffer bytes.Buffer
	var sigProofBuffer bytes.Buffer
	var chanllengeCollection []pocutil.Hash
	var proofCollection []Proof

	sk, pk, err := privateKeyFromString(sSk)
	if err != nil {
		return err
	}
	mdb, err := massdb_v1.NewMassDBV1(rootPath, engine.UnknownOrdinal, pk, bl)
	pkHash := mdb.PubKeyHash()

	// 对输入challenge反序列化
	bChallenges, err := hex.DecodeString(sChallenge)
	if err != nil {
		return err
	}
	challengeDecoder := gob.NewDecoder(bytes.NewReader(bChallenges))
	if err := challengeDecoder.Decode(&chanllengeCollection); err != nil {
		return err
	}

	for _, challenge := range chanllengeCollection {
		if proof, err := getProof(mdb, challenge, pkHash); err == nil {
			proofCollection = append(proofCollection, Proof{X: proof.X, Xp: proof.XPrime, Challenge: challenge, Success: true})
		} else {
			proofCollection = append(proofCollection, Proof{Challenge: challenge, Success: false})
		}
	}
	proofEncoder := gob.NewEncoder(&proofBuffer)
	if err := proofEncoder.Encode(&proofCollection); err != nil {
		fmt.Println("Failed to decode the proof\n")
		return err
	}
	proofHash := sha256.Sum256(proofBuffer.Bytes())
	signature, err := sk.Sign(proofHash[:])
	bSignature := signature.Serialize()
	sigProof := SigProof{ProofSer: proofBuffer.Bytes(), SigSer: bSignature, Bl: bl}

	sigEncoder := gob.NewEncoder(&sigProofBuffer)
	if err := sigEncoder.Encode(&sigProof); err != nil {
		fmt.Println("Failed to decode the proof")
		return err
	}
	fmt.Printf("Proof:\n%s\n", hex.EncodeToString(sigProofBuffer.Bytes()))

	return nil
}
